import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
import numpy as np
import os
class TactileDataset(Dataset):
    def __init__(self, tacdata, deltas, mean,std):
        """
        参数：
        tacdata : numpy数组 (num_samples, 400, 9)
        deltas : numpy数组 (num_samples,)
        """
        
        assert tacdata.shape[1:] == (400, 9), f"无效的触觉数据形状: {tacdata.shape}"
        assert len(deltas) == tacdata.shape[0], "数据与标签数量不匹配"
        
        # 在Dataset类中添加

        # 前处理
        tacdata = (tacdata - mean) / std

        self.tacdata = torch.tensor(tacdata, dtype=torch.float32)
        self.deltas = torch.tensor(deltas, dtype=torch.float32)
        self.augment = True  # 是否进行数据增强


        
        # 调整形状
        self.tacdata = self.tacdata.view(-1, 20, 20, 9)
        self.tacdata = self.tacdata.permute(0, 3, 1, 2)
        self.augment_transform = transforms.Compose([
            transforms.RandomHorizontalFlip(p=0.5),
            transforms.RandomVerticalFlip(p=0.5),
            transforms.RandomRotation(90)
        ])

        

    def __len__(self):
        return len(self.deltas)

    def __getitem__(self, idx):
        data = self.tacdata[idx]
        if self.augment:
            data = self.augment_transform(data)
        return data, self.deltas[idx]

class StreamTactileDataset(Dataset):
    """
    流式触觉数据集，适用于在线学习或增量学习场景
    """
    def __init__(self, metas, mean, std):
        """
        参数：
        meta : list, 每个元素是一个元组，(exp_dir, sensor_id, data_length)
        mean : numpy数组, 触觉数据的均值
        std : numpy数组, 触觉数据的标准差
        """
        
        self.metas = metas
        self.mean = mean
        self.std = std
        self.length = sum(meta[2] for meta in metas)
        self.index_prefix_sum = np.cumsum([0] + [meta[2] for meta in metas[:-1]])


        self.augment = True  # 是否进行数据增强

        self.augment_transform = transforms.Compose([
            transforms.RandomHorizontalFlip(p=0.5),
            transforms.RandomVerticalFlip(p=0.5),
            transforms.RandomRotation(90)
        ])

        

    def __len__(self):
        return self.length

    def __getitem__(self, idx):
        goal_meta_idx = np.searchsorted(self.index_prefix_sum, idx, side='right') - 1
        goal_meta = self.metas[goal_meta_idx]
        exp_dir, sensor_id, data_length = goal_meta
        local_idx = idx - self.index_prefix_sum[goal_meta_idx]
        
        deltas = np.load(os.path.join(exp_dir, f'true_delta_{sensor_id}.npy'))
        tacdata = np.load(os.path.join(exp_dir, f'tacdata_{sensor_id}.npy'))
        tacdata = torch.tensor(tacdata[local_idx]).view(20, 20, 9)
        deltas = torch.tensor(deltas[local_idx]).float()
        tacdata = ((tacdata - self.mean) / self.std).float().permute(2, 0, 1)
        if self.augment:
            tacdata = self.augment_transform(tacdata)
        return tacdata, deltas